# deadlock_env.py (V4 - Enhanced Complexity)
import copy
from typing import Tuple

import numpy as np
import random

from numpy import ndarray

# 定义进程状态
STATE_RUNNING = 0
STATE_BLOCKED = 1
STATE_TERMINATED = 2

KAPPA_PREEMPT = 0.3
KAPPA_ROLLBACK = 0.9
REWARD_SAFE = 1.0
REWARD_UNBLOCK = 0.3
PENALTY_STEP = -0.05
PENALTY_INVALID_ACTION = -0.1
REWARD_VALID_ACTION = 0.05

class DeadlockEnv:
    def __init__(self, num_processes: int, num_resources: int):
        # 为了能构造更复杂的场景，建议进程数大于资源数
        assert num_processes > num_resources, "建议进程数 > 资源数以构造更复杂的非循环阻塞"
        assert num_resources >= 3, "需要至少3个资源来构造循环"
        self.num_processes = num_processes
        self.num_resources = num_resources
        self.total_resources = None
        self.available = None
        self.allocation = None
        self.max_demand = None
        self.need = None
        self.processes_status = None
        self.process_properties = None
        self.unblocked_in_episode = set()

    def _initialize_state(self):
        """初始化一个干净的、无分配的状态。"""
        self.total_resources = np.array([random.randint(2, 4) for _ in range(self.num_resources)])
        self.available = self.total_resources.copy()
        self.processes_status = np.full(self.num_processes, STATE_RUNNING)
        self.process_properties = [{'is_realtime': random.choice([True, False])} for _ in range(self.num_processes)]
        self.allocation = np.zeros((self.num_processes, self.num_resources), dtype=int)
        self.max_demand = np.zeros((self.num_processes, self.num_resources), dtype=int)
        self.need = self.max_demand.copy()

    def _generate_randomized_deadlock(self):
        """
        [V2 - 强化版] 构造一个包含"主死锁环"和"次要独立死锁"的复杂场景。
        这确保了解决问题至少需要两步。
        """
        self._initialize_state()

        # --- 1. 构造主死锁环 (至少需要3个进程/资源) ---
        # 为了保证主环的复杂度，我们让它的大小至少为3
        main_cycle_min_size = 3
        # 主环最大不能超过总资源数的一半，为次要死锁留出空间
        main_cycle_max_size = max(main_cycle_min_size, self.num_resources // 2)

        main_cycle_size = random.randint(main_cycle_min_size, main_cycle_max_size)

        # 随机挑选进程和资源参与主循环
        p_in_main_cycle = sorted(random.sample(range(self.num_processes), main_cycle_size))
        r_in_main_cycle = sorted(random.sample(range(self.num_resources), main_cycle_size))

        for i in range(main_cycle_size):
            p_current_idx = p_in_main_cycle[i]
            r_hold_idx = r_in_main_cycle[i]
            r_request_idx = r_in_main_cycle[(i + 1) % main_cycle_size]

            # 分配资源
            if self.available[r_hold_idx] > 0:
                hold_amount = random.randint(1, self.available[r_hold_idx])
                self.allocation[p_current_idx, r_hold_idx] = hold_amount
                self.available[r_hold_idx] -= hold_amount
                self.max_demand[p_current_idx, r_hold_idx] = hold_amount

            # 设置需求并阻塞进程
            request_amount = random.randint(1, self.total_resources[r_request_idx])
            self.max_demand[p_current_idx, r_request_idx] = request_amount
            self.need[p_current_idx, r_request_idx] = request_amount
            self.processes_status[p_current_idx] = STATE_BLOCKED

        # --- 2. 构造次要独立死锁 (至少需要2个进程/资源) ---
        # 找出未被主环使用的进程和资源
        p_available = [p for p in range(self.num_processes) if p not in p_in_main_cycle]
        r_available = [r for r in range(self.num_resources) if r not in r_in_main_cycle]

        # 确保有足够的进程和资源来创建第二个死锁
        if len(p_available) >= 2 and len(r_available) >= 2:
            # 挑选2个进程和2个资源创建简单的 P_A -> R_A -> P_B -> R_B -> P_A 循环
            p_in_sec_cycle = sorted(random.sample(p_available, 2))
            r_in_sec_cycle = sorted(random.sample(r_available, 2))

            # 进程A持有资源A，请求资源B
            p_A, p_B = p_in_sec_cycle[0], p_in_sec_cycle[1]
            r_A, r_B = r_in_sec_cycle[0], r_in_sec_cycle[1]

            # 分配和请求 for p_A
            if self.available[r_A] > 0:
                hold_amount = random.randint(1, self.available[r_A])
                self.allocation[p_A, r_A] = hold_amount
                self.available[r_A] -= hold_amount
                self.max_demand[p_A, r_A] = hold_amount
            request_amount = random.randint(1, self.total_resources[r_B])
            self.max_demand[p_A, r_B] = request_amount
            self.need[p_A, r_B] = request_amount
            self.processes_status[p_A] = STATE_BLOCKED

            # 分配和请求 for p_B
            if self.available[r_B] > 0:
                hold_amount = random.randint(1, self.available[r_B])
                self.allocation[p_B, r_B] = hold_amount
                self.available[r_B] -= hold_amount
                self.max_demand[p_B, r_B] = hold_amount
            request_amount = random.randint(1, self.total_resources[r_A])
            self.max_demand[p_B, r_A] = request_amount
            self.need[p_B, r_A] = request_amount
            self.processes_status[p_B] = STATE_BLOCKED


    def reset(self) -> tuple[dict, ndarray]:
        self.unblocked_in_episode.clear()
        self._generate_randomized_deadlock()

        while self.is_safe():
            self._generate_randomized_deadlock()
        return self._get_observation(), self._get_action_mask()

    def is_safe(self) -> bool:
        work = self.available.copy()
        finish = np.array([st == STATE_TERMINATED for st in self.processes_status])

        num_p = self.num_processes
        while True:
            found_process = False
            for i in range(num_p):
                if not finish[i] and np.all(self.need[i] <= work):
                    work += self.allocation[i]
                    finish[i] = True
                    found_process = True
            if not found_process:
                break
        return np.all(finish)

    def _get_observation(self) -> dict:
        return {
            "total_resources": self.total_resources.copy(),
            "available": self.available.copy(),
            "allocation": self.allocation.copy(),
            "need": self.need.copy(),
            "processes_status": self.processes_status.copy(),
            "process_properties": copy.deepcopy(self.process_properties)
        }

    def render(self):
        print("=" * 50)
        blocked_pids = [i for i, s in enumerate(self.processes_status) if s == STATE_BLOCKED]
        print(f" 随机死锁场景 (涉及 {len(blocked_pids)} 个阻塞进程)")
        print("=" * 50)
        print(f"资源总量: {self.total_resources}")
        print(f"当前可用资源 (Available): {self.available}")
        print("\n--- 进程状态 ---")
        status_map = {0: "Running", 1: "Blocked", 2: "Terminated"}
        for i in range(self.num_processes):
            prop = "RT" if self.process_properties[i]['is_realtime'] else "NRT"
            print(f"  P{i:<2} | Status: {status_map[self.processes_status[i]]:<10} | Type: {prop}")
        print("\n--- 矩阵 ---")
        print("Allocation (已分配 A):")
        print(self.allocation)
        print("\nNeed (尚需 D-A):")
        print(self.need)
        print(f"\n当前状态是否安全? -> {self.is_safe()}")
        print("=" * 50)

    def _update_system_after_release(self) -> float:
        unblock_reward = 0.0
        while True:
            unlocked_in_this_pass = False
            for pid in range(self.num_processes):
                if self.processes_status[pid] == STATE_BLOCKED and np.all(self.need[pid] <= self.available):
                    self.available += self.allocation[pid]
                    self.allocation[pid, :] = 0
                    self.need[pid, :] = 0
                    self.processes_status[pid] = STATE_TERMINATED
                    if pid not in self.unblocked_in_episode:
                        unblock_reward += REWARD_UNBLOCK
                        self.unblocked_in_episode.add(pid)
                    unlocked_in_this_pass = True
            if not unlocked_in_this_pass:
                break
        return unblock_reward

    def step(self, action: dict) -> tuple:
        action_type = action.get('type')
        pid = action.get('process_id')
        if pid is None or not (0 <= pid < self.num_processes):
            info = {'action_cost': 0.0, 'valid': False, 'reason': 'Invalid Process ID'}
            reward = PENALTY_INVALID_ACTION
            obs = self._get_observation()
            # No print statement here, as this is a fundamental error
            return obs, reward, False, info, self._get_action_mask()
        reward = PENALTY_STEP
        done = False
        info = {'action_cost': 0.0, 'valid': True, 'reason': 'Success'}

        has_resources = np.sum(self.allocation[pid]) > 0

        if action_type == 'preempt':
            is_realtime = self.process_properties[pid]['is_realtime']
            if not is_realtime:
                info.update({'valid': False, 'reason': 'Attempted to preempt a non-real-time process'})
            elif not has_resources:
                info.update({'valid': False, 'reason': 'Attempted to preempt a process with no resources'})
            else:
                self.available += self.allocation[pid]
                self.allocation[pid, :] = 0
                self.need[pid, :] = self.max_demand[pid, :]
                self.processes_status[pid] = STATE_RUNNING
                info['action_cost'] = KAPPA_PREEMPT

        elif action_type == 'rollback':
            if not has_resources:
                info.update({'valid': False, 'reason': 'Attempted to rollback a process with no resources'})
            else:
                self.available += self.allocation[pid]
                self.allocation[pid, :] = 0
                self.need[pid, :] = self.max_demand[pid, :]
                self.processes_status[pid] = STATE_RUNNING
                info['action_cost'] = KAPPA_ROLLBACK
        else:
            info.update({'valid': False, 'reason': f"Unknown action type: {action_type}"})

        if not info['valid']:
            # print(f"Invalid action on P{pid}: {info['reason']}")
            reward += PENALTY_INVALID_ACTION
            obs = self._get_observation()
            return obs, reward, done, info, self._get_action_mask()

        # If the action was valid:
        reward += REWARD_VALID_ACTION
        reward -= info['action_cost']

        # Check if the action unblocked any processes
        unblock_reward = self._update_system_after_release()
        reward += unblock_reward

        # Check if the system is now safe (deadlock resolved)
        if self.is_safe():
            done = True
            reward += REWARD_SAFE

        obs = self._get_observation()
        return obs, reward, done, info, self._get_action_mask()


    def _get_action_mask(self) -> np.ndarray:
        """
        [修改后]
        计算并返回当前状态下所有动作的有效性掩码。
        一个动作是有效的，当且仅当它能对环境状态产生实际影响。
        """
        num_actions = 2 * self.num_processes
        mask = np.zeros(num_actions, dtype=bool)

        for pid in range(self.num_processes):
            # 基础条件：进程必须尚未终止，并且实际持有资源
            if self.processes_status[pid] != STATE_TERMINATED and np.sum(self.allocation[pid]) > 0:
                # 1. Rollback 动作的有效性
                # 只要满足基础条件，就可以执行 rollback
                rollback_action_id = self.num_processes + pid
                mask[rollback_action_id] = True

                # 2. Preempt 动作的有效性
                # 只有当进程是实时进程(realtime)时，才能执行 preempt
                if self.process_properties[pid]['is_realtime']:
                    preempt_action_id = pid
                    mask[preempt_action_id] = True
        return mask